#!/usr/bin/python

#ass
import os,logging,scipy,math 
import random

logging.basicConfig(filename='atom_gen.log', level=logging.INFO)

# Total number of atoms 
num_of_atoms = 10000
# Degree of polymerization
degree_of_polym = 1000
# Bond length
bond_length = 1.53
# Infinitesimal
epsilon = 1e-5

# Lists atom positions
molecular_pos_x=[]
molecular_pos_y=[]
molecular_pos_z=[]

# Flag variable used by the loop
pass_flag = False

logging.info('Number of atoms: %d' % num_of_atoms)
logging.info('Degree of polymerization: %d' % degree_of_polym)
logging.info('Bond length: %d' % degree_of_polym)

# Loop into each atom [molecular_pos_x[ii],molecular_pos_y[ii],molecular_pos_z[ii]]
logging.info('Start looping into the atoms' + os.linesep)
for ii in xrange(num_of_atoms): 
  
    pass_flag = False
 
    while not pass_flag:
                
        if ii % degree_of_polym != 0:
            angle_xy = scipy.rand()*math.pi*2
            angle_yz = scipy.rand()*math.pi*2
            angle_xz = scipy.rand()*math.pi*2
            
        
            x_cord = molecular_pos_x[ii-1] + bond_length * math.sin(angle_yz)
            y_cord = molecular_pos_y[ii-1] + bond_length * math.sin(angle_xz)
            z_cord = molecular_pos_z[ii-1] + bond_length * math.sin(angle_xy)
        else:
            x_cord = random.uniform(0,80)
            y_cord = random.uniform(0,80)
            z_cord = random.uniform(0,80)
            
        
        for jj in xrange(ii):
            if ((molecular_pos_x[jj] - x_cord)**2 +  (molecular_pos_y[jj] 
                - y_cord)**2 + (molecular_pos_z[jj] - z_cord)**2) <= epsilon **2 : 
                
                pass_flag = False
                break
            
            if (jj % degree_of_polym != 0) and ((x_cord > (80 - epsilon/2) 
                    or x_cord < (0 + epsilon/2)) or (y_cord > (80 - epsilon/2) 
                    or y_cord < (0 + epsilon/2)) or (z_cord > (80 - epsilon/2) 
                    or z_cord < (0 + epsilon/2))): 
                pass_flag = False
                break
        else:
            pass_flag = True

  
    molecular_pos_x.append(x_cord)
    molecular_pos_y.append(y_cord)
    molecular_pos_z.append(z_cord)

# Start a new data file

logging.info('Finished atom positioning, saving to file...')
with open('data.dat', 'w') as datafile:
    datafile.write('# Model for PE' + os.linesep)
    datafile.write(os.linesep)
    datafile.write('      %d     atoms' % num_of_atoms + os.linesep)
    datafile.write('      %d     bonds' % 
            (num_of_atoms - int(num_of_atoms / degree_of_polym)) + os.linesep)
    datafile.write('      %d     angles' % 
            (num_of_atoms - int(num_of_atoms / degree_of_polym) * 2) + os.linesep)
    datafile.write('      %d     dihedrals' % 
            (num_of_atoms - int(num_of_atoms / degree_of_polym) * 3) + os.linesep)
    datafile.write(os.linesep)
    datafile.write(' + os.linesep)        1     atom types' + os.linesep)
    datafile.write('         1     bond types' + os.linesep)
    datafile.write('         1     angle types' + os.linesep)
    datafile.write('         1     dihedral types' + os.linesep)
    datafile.write(os.linesep)
    datafile.write('    0.0000   80.0586 xlo xhi' + os.linesep)
    datafile.write('    0.0000   80.0586 ylo yhi' + os.linesep)
    datafile.write('    0.0000   80.0586 zlo zhi' + os.linesep)
    datafile.write(os.linesep)
    datafile.write('Masses' + os.linesep)
    datafile.write(os.linesep)
    datafile.write('    1    14.02' + os.linesep)
    datafile.write(os.linesep)

import csv
with open('data.dat','a+') as csvfile:
    writer = csv.writer(csvfile,delimiter='\t',
            quotechar='',quoting=csv.QUOTE_NONE, dialect='excel')

    # write to the atoms
    logging.info('Write to the atoms...')
    csvfile.write(os.linesep)
    csvfile.write('Atoms' + os.linesep)
    csvfile.write(os.linesep)
    for ii in range(len(molecular_pos_x)): 
        writer.writerow(
            ['', ii+1, 1, 1, molecular_pos_x[ii],molecular_pos_y[ii],molecular_pos_z[ii]]
            )

    # write to the bonds
    logging.info('Write to the bonds...')
    csvfile.write(os.linesep)
    csvfile.write('Bonds' + os.linesep)
    csvfile.write(os.linesep)
    bond_num = 1
    for ii in range(len(molecular_pos_x)): 
        if ii % degree_of_polym != degree_of_polym - 1: 
            writer.writerow(
                ['', bond_num, 1, ii+1, ii+2]
                )
            bond_num += 1

    # write to angles
    logging.info('Write to the angles...')
    csvfile.write(os.linesep)
    csvfile.write('Angles' + os.linesep)
    csvfile.write(os.linesep)
    angle_num = 1
    for ii in range(len(molecular_pos_x)): 
        if ii % degree_of_polym < degree_of_polym - 2: 
            writer.writerow(
                ['', angle_num, 1, ii+1, ii+2, ii+3]
                )
            angle_num += 1


    # write to dihedrals 
    logging.info('Write to the angles...')
    csvfile.write(os.linesep)
    csvfile.write('Dihedrals' + os.linesep)
    csvfile.write(os.linesep)
    dihedral_num = 1
    for ii in range(len(molecular_pos_x)): 
        if ii % degree_of_polym < degree_of_polym - 3: 
            writer.writerow(
                ['', dihedral_num, 1, ii+1, ii+2, ii+3, ii+4]
                )
            dihedral_num += 1


